import pypose as pp
import torch


class CartPole(pp.module.NLS):
    def __init__(self, dt, length, cartmass, polemass, gravity):
        super().__init__()
        self.tau = dt
        self.length = length
        self.cartmass = cartmass
        self.polemass = polemass
        self.gravity = gravity
        self.polemassLength = self.polemass * self.length
        self.totalMass = self.cartmass + self.polemass

    def state_transition(self, state, input, t=None):
        x = state[:, 0]
        xDot = state[:, 1]
        theta = state[:, 2]
        thetaDot = state[:, 3]

        force = input.squeeze()
        costheta = theta.cos()
        sintheta = theta.sin()

        temp = (force + self.polemassLength * thetaDot**2 * sintheta) / self.totalMass

        thetaAcc = (self.gravity * sintheta - costheta * temp) / (
            self.length * (4.0 / 3.0 - self.polemass * costheta**2 / self.totalMass)
        )

        xAcc = temp - self.polemassLength * thetaAcc * costheta / self.totalMass

        _dstate = torch.stack((xDot, xAcc, thetaDot, thetaAcc), dim=-1)

        return state + _dstate * self.tau

    def observation(self, state, input, t=None):
        return state
